import os
import re
import traceback
from openpyxl.reader.excel import load_workbook

pattern = r"[`']([^`']+)[`']"
tradition_ = 'COMMENT'
common_part = 'STRING COMMENT '
last_line_str = ' PARTITIONED BY (As BIGINT)\nSTORED AS PARQUET;'


def modify_sheet(old_sheet, new_sheet_):
    constant_space_num = 30
    for row in old_sheet.iter_rows(values_only=True):
        origin_sql = row[0]
        if 'CREATE TABLE' in origin_sql:
            new_sheet_.append([origin_sql])
            continue
        if 'COMMENT=' in origin_sql:
            # 最后一行语句
            comment = origin_sql.split('COMMENT=')[-1].replace(';', '').replace(' ', '')
            new_sql = ')' + '\n' + 'COMMENT ' + comment + '\n' + last_line_str
            new_sheet_.append([new_sql])
            break
        cut_result: list = re.findall(pattern, origin_sql)
        if len(cut_result) > 1 and tradition_ in origin_sql:
            field, comment = cut_result[0], cut_result[-1]
            new_sql = field.ljust(constant_space_num) + common_part + f"'{comment}'" + ','
            new_sheet_.append([new_sql])


current_dir = os.getcwd()

try:
    for file in os.listdir(current_dir):
        if file.endswith(".xlsx") or file.endswith(".xls"):
            file_path: str = os.path.join(current_dir, file)
            book = load_workbook(file_path)
            # sheet0是你的原始数据,创建一个新的sheet保存修改后的数据
            origin_sheet, new_sheet = book.worksheets[0], book.create_sheet()
            modify_sheet(origin_sheet, new_sheet)
            book.save(file_path)

except:
    print(traceback.format_exc())
